package com.ty.utils;

import boofcv.alg.misc.ImageStatistics;
import boofcv.io.image.ConvertBufferedImage;
import boofcv.struct.image.GrayF32;
import org.opencv.core.Mat;
import org.opencv.imgproc.Imgproc;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

/**
 * 图片工具类
 *
 * @Author Tommy
 * @Date 2024/4/19
 */
public class ImageUtil {

    /**
     * 读取图片文件
     *
     * @param imgFile 图片文件对象
     * @return BufferedImage
     */
    public static BufferedImage read(File imgFile) throws IOException {
        return ImageIO.read(imgFile);
    }

    /**
     * 两张图片的尺寸是否一致
     *
     * @param image1 图片对象1
     * @param image2 图片对象2
     * @return boolean
     */
    public static boolean sameSize(BufferedImage image1, BufferedImage image2) {
        if (null == image1 || null == image2) {
            return false;
        }
        return image1.getWidth() == image2.getWidth() && image1.getHeight() == image2.getHeight();
    }

    /**
     * 转换为灰度图
     *
     * @param image 图片对象
     * @return GrayF32
     */
    public static GrayF32 toGray(BufferedImage image) {
        return ConvertBufferedImage.convertFrom(image, (GrayF32) null);
    }

    /**
     * 计算灰度图方差
     *
     * @param gray 灰度图对象
     * @return float
     */
    public static float variance(GrayF32 gray) {
        float mean = ImageStatistics.mean(gray);
        return ImageStatistics.variance(gray, mean);
    }

    /**
     * 计算两张灰度图的MSE（Mean Square Error）
     *
     * @param gray1 灰度图对象1
     * @param gray2 灰度图对象2
     * @return double
     */
    public static double mse(GrayF32 gray1, GrayF32 gray2) {
        return ImageStatistics.meanDiffSq(gray1, gray2);
    }

    /**
     * 计算两张灰度图的NMSE（Normalized Mean Square Error）
     *
     * @param gray1 灰度图对象1
     * @param gray2 灰度图对象2
     * @return double
     */
    public static double nmse(GrayF32 gray1, GrayF32 gray2) {
        double mse = mse(gray1, gray2);
        float variance = variance(gray1);
        return mse / variance;
    }

    /**
     * 计算两张灰度图的相似度
     *
     * @param gray1 灰度图对象1
     * @param gray2 灰度图对象2
     * @return double 返回相似度的值，范围：[0, 1]
     */
    public static double similarity(GrayF32 gray1, GrayF32 gray2) {
        double similarity = 0d;
        double nmse = nmse(gray1, gray2);
        if (nmse >= 0 && nmse <= 1) {
            similarity = Math.round((1 - nmse) * 10000d);
            similarity /= 10000d;
        }
        return similarity;
    }

    /**
     * 图片是否有缺损
     *
     * @param gray      灰度图对象
     * @param threshold 判定为缺损的阈值比例
     * @return boolean 返回True / False
     */
    public static boolean isLack(GrayF32 gray, double threshold) {
        // 按行遍历图片，若连续行数的像素值均一致，且行数占比超过 图片比例 threshold，则认为图片是缺损的
        int maxLossLines = (int)Math.floor(gray.getHeight() * threshold);

        boolean flag = false;
        int tick = 0;
        for (int y = 0; y < gray.getHeight(); y++) {
            Set<Float> lineVSet = new HashSet<>();
            for (int x = 0; x < gray.getWidth(); x++) {
                float v = gray.unsafe_get(x, y);
                lineVSet.add(v);
            }

            if (1 == lineVSet.size()) {
                tick++;
                if (tick > maxLossLines) { // 图片是缺损的
                    flag = true;
                    break;
                }
            } else {
                tick = 0;
            }
        }
        return flag;
    }

    /**
     * OpenCV Mat 转 BufferedImage
     */
    public static BufferedImage mat2Image(Mat mat) {
        int width = mat.width();
        int height = mat.height();
        byte[] data = new byte[width * height * (int) mat.elemSize()];
        Imgproc.cvtColor(mat, mat, Imgproc.COLOR_BGR2RGB);
        mat.get(0, 0, data);
        BufferedImage ret = new BufferedImage(width, height, BufferedImage.TYPE_3BYTE_BGR);
        ret.getRaster().setDataElements(0, 0, width, height, data);
        return ret;
    }
}
